"""
连通域算法实现
"""

import numpy as np
import util.img_utils as img_utils

def img_connected_domain(img_np, threshold):
    is_sum = len(img_np.shape) >= 3
    height, width = img_np.shape[0], img_np.shape[1]

    # 标签集合
    label_list = []
    label_list_index = 0
    all_label = np.zeros((height, width), dtype=int)
    label_list.append(set([0]))

    # 一轮打标签
    for x in range(0, height):
        for y in range(0, width):
            if is_sum:
                if sum(img_np[x, y]) <= threshold:
                    label_list_index = set_label(x, y, all_label, label_list, label_list_index)
            else:
                if img_np[x, y] <= threshold:
                    label_list_index = set_label(x, y, all_label, label_list, label_list_index)

    # 简化标签关系
    simple_label_list = {}
    # 所有标签
    tmp_label = []
    for s in label_list:
        min_label = min(s)
        if min_label != 0:
            tmp_label.append(min_label)
        for j in s:
            simple_label_list[j] = min_label
    label_list = simple_label_list
    # print(tmp_label)

    # 二轮同步标签
    for x in range(0, height):
        for y in range(0, width):
            all_label[x][y] = label_list[all_label[x][y]]


    # 统计有多少个标签并给颜色
    # label_color = {}
    # if is_sum:
    #     static_color = [(0, 0, 0), (0, 0, 128), (0, 255, 255), (0, 100, 0), (85, 107, 47), (0, 255, 0), (34, 139, 34),
    #                     (255, 255, 0)
    #         , (205, 92, 92), (210, 105, 30)]
    # else:
    #     static_color = [0, 40, 80, 120, 160, 200]
    # now_color_index = 0
    # for i in label_list:
    #     if label_list[i] == 0:
    #         label_color[i] = 255
    #     if label_list[i] not in label_color:
    #         label_color[i] = static_color[now_color_index % len(static_color)]
    #         now_color_index = now_color_index + 1
    # label_color.pop(0)
    # print(label_color, len(label_color))

    all_img = {}
    all_img_index = []
    for one_label in tmp_label:
        ima, min_w = cut_by_label(img_np, one_label, all_label, threshold)
        if ima is not None:
            while min_w in all_img:
                min_w = min_w + 1
            all_img[min_w] = ima
            all_img_index.append(min_w)

    all_img_index.sort()
    all_img_sort = []
    for i in all_img_index:
        all_img_sort.append(all_img[i])

    # 三轮给颜色
    # for x in range(0, height):
    #     for y in range(0, width):
    #         if all_label[x][y] != 0:
    #             img_np[x][y] = label_color[all_label[x][y]]

    return all_img_sort


def set_label(height, width, all_label, label_list, label_list_index):
    if all_label[height][width] != 0:
        return
    # 上面一位的标签
    above_label = None
    if height != 0 and all_label[height - 1][width] != 0:
        above_label = all_label[height - 1][width]
    # 前面一位的标签
    ahead_label = None
    if width != 0 and all_label[height][width - 1] != 0:
        ahead_label = all_label[height][width - 1]

    if above_label is None:
        if ahead_label is None:
            label_list_index = label_list_index + 1
            new_label = label_list_index
            new_label_set = set([new_label])
            label_list.append(new_label_set)
            all_label[height][width] = new_label
        else:
            all_label[height][width] = ahead_label
    else:
        if ahead_label is None:
            all_label[height][width] = above_label
        else:
            if above_label == ahead_label:
                all_label[height][width] = above_label
            else:
                # 选择最小值
                min_label = min(above_label, ahead_label)
                all_label[height][width] = min_label

                # 合并关系
                above_label_set = pop_label_set(above_label, label_list)
                if ahead_label not in above_label_set:
                    ahead_label_index = pop_label_set(ahead_label, label_list)
                    one_set = set.union(above_label_set, ahead_label_index)
                else:
                    one_set = above_label_set
                label_list.append(one_set)

    return label_list_index


def pop_label_set(label, label_list):
    for i in range(len(label_list)):
        if label in label_list[i]:
            return label_list.pop(i)


def cut_by_label(img_np, label, all_label, threshold):
    w_max = np.argmax(all_label == label, axis=0)

    min_w, max_w = None, None
    for i in range(w_max.shape[0]):
        if w_max[i]:
            if min_w is None:
                min_w = i
                max_w = i
            else:
                max_w = i

    if min_w is None:
        return None, 0
    # print(min_w, max_w, max_w - min_w)
    # 宽度太小的不要
    if max_w - min_w < 3:
        return None, 0
    img_np = img_np.copy()
    img_np[all_label != label] = 255
    # 切掉白边
    img_np = img_utils.cut_white_area(img_np, threshold)
    return img_np, min_w
